import time
import os
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, VotingClassifier, BaggingClassifier, AdaBoostClassifier, StackingClassifier
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from xgboost import XGBClassifier
from tabulate import tabulate
import matplotlib.pyplot as plt
from tqdm import tqdm
import logging  # 导入日志记录库
import joblib  # 用于保存和加载模型
from util import createXY  # 假设createXY函数已经在util.py中实现

# 配置日志格式
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 使用createXY函数从data文件夹中加载图像数据
train_folder = "./data/train"
dest_folder = "./data"
X, y = createXY(train_folder, dest_folder, method='flat')

# 转换为numpy数组
X = np.array(X)
y = np.array(y)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)

# 调整后的逻辑回归分类器
logistic_clf = LogisticRegression(C=0.5, penalty='l2', solver='lbfgs', max_iter=1000)

# 调整后的随机森林分类器
random_forest_clf_subcls = RandomForestClassifier(n_estimators=500, max_depth=10, min_samples_split=5, min_samples_leaf=2, random_state=42)

# 调整后的支持向量机分类器
svm_clf = SVC(C=1, gamma='scale', probability=True)

# 创建硬投票分类器，包含逻辑回归、随机森林和支持向量机三个分类器
voting_clf_hard = VotingClassifier(
    estimators=[('lr', logistic_clf), ('rf', random_forest_clf_subcls), ('svc', svm_clf)],
    voting='hard', # 采用硬投票, 即选择票数最多的类别作为预测结果
)

# 创建软投票分类器，包含逻辑回归、随机森林和支持向量机三个分类器
voting_clf_soft = VotingClassifier(
    estimators=[('lr', logistic_clf), ('rf', random_forest_clf_subcls), ('svc', svm_clf)],
    voting='soft', # 采用软投票, 即选择概率之和最大的类别作为预测结果
    weights=[1, 1, 2],  # 为支持向量机分类器的权重设置为2，其他分类器的权重为1
)

# 调整后的Bagging分类器
bag_clf = BaggingClassifier(
    DecisionTreeClassifier(max_depth=10), # 增加基分类器的深度
    n_estimators=500,                    # 增加基分类器的数量
    max_samples=200,                      # 增加每个基分类器的样本数量
    bootstrap=True,
    random_state=42
)

# 创建Pasting分类器
paste_clf = BaggingClassifier(
    DecisionTreeClassifier(), # 基分类器选用决策树分类器
    n_estimators=500,         # 500个分类器
    max_samples=100,          # 每个分类器的训练集包含100个样本
    bootstrap=False,          # 无放回的采样
    random_state=42           # 随机种子
)

# 创建有线程优化的Bagging分类器
bag_clf_multicore = BaggingClassifier(
    DecisionTreeClassifier(), # 基分类器选用决策树分类器
    n_estimators=500,         # 500个分类器
    max_samples=100,          # 每个分类器的训练集包含100个样本
    bootstrap=True,           # 有放回的采样
    n_jobs=-1,                # 使用所有CPU核
    random_state=42           # 随机种子
)

# 创建有线程优化的决策树分类器
random_forest_clf = RandomForestClassifier(
    n_estimators=500,  # 500个分类器
    max_samples=100,   # 每个分类器的训练集包含100个样本
    bootstrap=True,    # 有放回的采样
    n_jobs=-1,         # 使用所有CPU核
    random_state=42    # 随机种子
)

# 创建有线程优化的Pasting分类器
paste_clf_multicore = BaggingClassifier(
    DecisionTreeClassifier(), # 基分类器选用决策树分类器
    n_estimators=500,         # 500个分类器
    max_samples=100,          # 每个分类器的训练集包含100个样本
    bootstrap=False,          # 无放回的采样
    n_jobs=-1,                # 使用所有CPU核
    random_state=42           # 随机种子
)

# 调整后的AdaBoost分类器
ada_clf = AdaBoostClassifier(
    DecisionTreeClassifier(max_depth=2),  # 增加基分类器的深度
    n_estimators=300,                     # 增加弱分类器的数量
    learning_rate=0.1                     # 降低学习率
) 

# 调整后的XGBoost分类器
xgb_clf_v1 = XGBClassifier(
    n_estimators=300,       # 增加树的数量
    learning_rate=0.05,     # 降低学习率
    max_depth=6,            # 增加树的深度
    min_child_weight=3,     # 调整最小子节点样本权重和
    subsample=0.9,          # 增加每棵树随机采样的样本比例
    colsample_bytree=0.9,   # 增加每棵树随机采样的特征比例
    gamma=0.2,              # 调整节点分裂所需的最小损失函数下降值
    random_state=42
)

xgb_clf_v2 = XGBClassifier(
    n_estimators=150,       # 增加树的数量
    learning_rate=0.05,     # 降低学习率
    max_depth=4,            # 增加树的深度
    min_child_weight=1,     # 调整最小子节点样本权重和
    subsample=0.8,          # 增加每棵树随机采样的样本比例
    colsample_bytree=0.8,   # 增加每棵树随机采样的特征比例
    gamma=0.1,              # 调整节点分裂所需的最小损失函数下降值
    random_state=42
)

# 创建一个堆叠分类器
stacking_clf = StackingClassifier(
    estimators=[('lr', logistic_clf), ('rf', random_forest_clf_subcls), ('svm', svm_clf)],  # 基础模型
    final_estimator=LogisticRegression(),  # 元学习器
    cv=5,  # 交叉验证次数
    stack_method='auto'  # 堆叠方式，可以是 'predict_proba', 'decision_function' 或 'predict'
)

# 更新分类器字典
clfs = {
    "logistic_regression": logistic_clf,
    "random_forest_subcls": random_forest_clf_subcls,
    "svm": svm_clf,
    "hard_voting": voting_clf_hard,
    "soft_voting": voting_clf_soft,
    "bagging": bag_clf,
    "bagging_multicore": bag_clf_multicore,
    "random_forest": random_forest_clf,
    "pasting": paste_clf,
    "pasting_multicore": paste_clf_multicore,
    "adaboost": ada_clf,
    "gradient_boosting_v1": xgb_clf_v1,
    "gradient_boosting_v2": xgb_clf_v2,
    "stacking": stacking_clf
}

results = []
# 训练和评估分类器
for name, clf in tqdm(clfs.items()):# 遍历字典中的每一个键值对:(模型名称, 模型对象)
    logging.info(f"开始训练模型: {name}")
    start_time = time.time()
    clf.fit(X_train, y_train)
    fit_time = time.time() - start_time
    logging.info(f"模型 {name} 训练完成，耗时: {fit_time:.2f} 秒")

    start_time = time.time()
    accuracy = clf.score(X_test, y_test)
    score_time = time.time() - start_time
    logging.info(f"模型 {name} 评估完成，耗时: {score_time:.2f} 秒，准确率: {accuracy:.4f}")

    results.append([name, fit_time, score_time, accuracy])

    # 保存模型
    model_filename = f"{name}_model.pkl"
    joblib.dump(clf, model_filename)
    logging.info(f"模型 {name} 已保存为 {model_filename}")

# 打印结果表格
headers = ["Classifier", "Training Time (s)", "Prediction Time (s)", "Accuracy"]
print(tabulate(results, headers=headers, tablefmt="simple"))